from .base import (
    PoseEstimator,
    center_ellipse,
    ellipse_args,
    ransac_pnp,
    pose_calculate,
)
import numpy as np
import cv2
from scipy.optimize import least_squares


class EP(PoseEstimator):
    def __init__(self, K, use_pnp_initial=True, *args, **kwargs):
        super().__init__(K, *args, **kwargs)
        self.use_pnp_initial = use_pnp_initial

    def __name__(self):
        return "ep"

    def forward(self, C_3d, C_2d,*args, **ignore_kwargs):
        if self.use_pnp_initial:
            x_cnt = []
            X_cnt = []
            for c_3d, c_2d in zip(C_3d, C_2d):
                # 求解相机椭圆中心：
                x_cnt.append(center_ellipse(c_2d))
                X_cnt.append(center_ellipse(c_3d))
            # 求解相机位姿
            X_cnt = np.pad(X_cnt, ((0, 0), (0, 1)), "constant", constant_values=0)
            x_cnt = np.array(x_cnt)
            ind = ransac_pnp(self.K, X_cnt, x_cnt)
            if ind is None:
                return False, None, None
            R, T = pose_calculate(self.K, X_cnt[ind], x_cnt[ind])
            initial_pose = np.concatenate((cv2.Rodrigues(R)[0], T)).flatten()
        else:
            initial_pose = np.zeros(6)
        C_2d = np.array(C_2d)
        C_3d = np.array(C_3d)
        return self.estimator(C_3d, C_2d, initial_pose)

    def estimator(self, database, measurements, initial_pose):
        """
        K: 相机内参
        database: 三维点库
        measurements: 观测到的二维点
        """
        assert len(database) == len(measurements)
        assert measurements.shape[1] == 3
        assert database.shape[1] == 3

        def residual_function(params, C_3d, C_2d):
            # 定义残差函数
            R = cv2.Rodrigues(params[:3])[0]
            t = params[3:]
            P = self.K @ np.hstack((R, t[:, None]))
            cone = P.T @ C_2d @ P
            cone = np.pad(cone, ((0, 0), (0, 1), (0, 1)), "constant", constant_values=0)
            C_esti = self.cone_plane_cross(cone, (0, 0, 1, 0))
            err = []
            for i in range(len(C_2d)):
                arg_real = ellipse_args(C_3d[i])
                arg_esti = ellipse_args(C_esti[i])
                if not arg_esti[0]:
                    err.append(np.zeros(5))
                else:
                    err.append(np.array(arg_real) - np.array(arg_esti))
            err = np.concatenate(err)
            err[np.isnan(err)] = 0
            return err

        result = least_squares(
            residual_function, initial_pose, args=(database, measurements)
        )
        if not result.success:
            return False, None, None
        t = result.x[3:]  # 提取平移部分
        R = cv2.Rodrigues(result.x[:3])[0]
        return True, R, t

    def cone_plane_cross(self, Q, plane):
        """
        求解圆锥与平面的交点
        cone: 圆锥
        plane: 平面
        """
        # 提取平面方程的参数
        a, b, c, d = plane

        # 计算交线的齐次矩阵 Q'
        A_prime = Q[:, 0, 0] * c**2 - 2 * a * c * Q[:, 0, 2] + (a**2) * Q[:, 2, 2]
        E_prime = Q[:, 1, 1] * c**2 - 2 * b * c * Q[:, 1, 2] + (b**2) * Q[:, 2, 2]
        B_prime = (
            Q[:, 0, 1] * c**2
            - a * c * Q[:, 1, 2]
            - b * c * Q[:, 0, 2]
            + (a * b) * Q[:, 2, 2]
        )
        D_prime = (
            Q[:, 0, 3] * c**2
            - a * c * Q[:, 2, 3]
            - d * c * Q[:, 0, 2]
            + (a * d) * Q[:, 2, 2]
        )
        G_prime = (
            Q[:, 1, 3] * c**2
            - b * c * Q[:, 2, 3]
            - d * c * Q[:, 1, 2]
            + (b * d) * Q[:, 2, 2]
        )
        J_prime = Q[:, 3, 3] * c**2 - 2 * d * c * Q[:, 2, 3] + (d**2) * Q[:, 2, 2]

        return np.array(
            [
                [A_prime, B_prime, D_prime],
                [B_prime, E_prime, G_prime],
                [D_prime, G_prime, J_prime],
            ]
        ).T
